from collections import OrderedDict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.conv import _ConvNd
from torch.nn.modules.utils import _pair

from networks.Fuck import base_upsample, ResNet18WithoutMaxPool, ResNet34WithoutMaxPool, RepVGGBlock, \
    ResNet10WithoutMaxPool


class GatedConv2d(_ConvNd):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=False):
        super(GatedConv2d, self).__init__(in_channels, out_channels,
                                          _pair(kernel_size),
                                          _pair(stride),
                                          _pair(padding),
                                          _pair(dilation),
                                          False, _pair(0), groups, bias, 'zeros')

        self.gate_conv = nn.Sequential(
            nn.BatchNorm2d(in_channels + 1),
            nn.Conv2d(in_channels + 1, in_channels + 1, 1),
            nn.ReLU(True),
            nn.Conv2d(in_channels + 1, 1, 1),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

    def forward(self, input_features, gating_features):
        """
        :param input_features:  [NxCxHxW]
        :param gating_features: [Nx1xHxW]
        :return:
        """
        alphas = self.gate_conv(torch.cat([input_features, gating_features], dim=1))
        input_features = (input_features * (alphas + 1))
        return F.conv2d(input_features, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)


class GatedConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1,
                 rep=False, deploy=False):
        super(GatedConvBNReLU, self).__init__()

        conv = RepConvBNReLU if rep else BaseConvBNReLU

        self.gate_conv = nn.Sequential(
            BaseConvBNReLU(in_channels + 1, in_channels + 1, 1, deploy=deploy),
            nn.Conv2d(in_channels + 1, 1, 1),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )
        self.out = conv(in_channels, out_channels, kernel_size, stride, padding, dilation, deploy=deploy)

    def forward(self, x, edge):
        alphas = self.gate_conv(torch.cat([x, edge], dim=1))
        return self.out(x * alphas + x)


class BaseConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, deploy=False):
        super(BaseConvBNReLU, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels,
                      out_channels,
                      kernel_size=kernel_size,
                      stride=stride,
                      padding=padding,
                      dilation=dilation,
                      bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class RepConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, deploy=False):
        super(RepConvBNReLU, self).__init__()
        self.conv = nn.Sequential(
            RepVGGBlock(in_channels,
                        out_channels,
                        kernel_size=kernel_size,
                        stride=stride,
                        padding=padding,
                        dilation=dilation,
                        deploy=deploy)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class RepConv1x1(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, deploy=False):
        super(RepConv1x1, self).__init__()

        self.relu = nn.ReLU(True)

        if deploy:
            self.rbr_reparam = nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=True, padding_mode='zeros')
        else:
            self.rbr_dense = nn.Sequential(OrderedDict([
                ('conv', nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False)),
                ('bn', nn.BatchNorm2d(num_features=out_channels))
            ]))

    def forward(self, x):
        if hasattr(self, 'rbr_reparam'):
            return self.relu(self.rbr_reparam(x))
        return self.relu(self.rbr_dense(x))

    def get_equivalent_kernel_bias(self):
        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_dense)
        return kernel1x1, bias1x1

    def _fuse_bn_tensor(self, branch):
        if branch is None:
            return 0, 0
        if isinstance(branch, nn.Sequential):
            kernel = branch.conv.weight
            running_mean = branch.bn.running_mean
            running_var = branch.bn.running_var
            gamma = branch.bn.weight
            beta = branch.bn.bias
            eps = branch.bn.eps
        else:
            assert isinstance(branch, nn.BatchNorm2d)
            if not hasattr(self, 'id_tensor'):
                input_dim = self.in_channels // self.groups
                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
                for i in range(self.in_channels):
                    kernel_value[i, i % input_dim, 1, 1] = 1
                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
            kernel = self.id_tensor
            running_mean = branch.running_mean
            running_var = branch.running_var
            gamma = branch.weight
            beta = branch.bias
            eps = branch.eps
        std = (running_var + eps).sqrt()
        t = (gamma / std).reshape(-1, 1, 1, 1)
        return kernel * t, beta - running_mean * gamma / std

    def switch_to_deploy(self):
        if hasattr(self, 'rbr_reparam'):
            return
        kernel, bias = self.get_equivalent_kernel_bias()
        self.rbr_reparam = nn.Conv2d(
            in_channels=self.rbr_dense.conv.in_channels,
            out_channels=self.rbr_dense.conv.out_channels,
            kernel_size=self.rbr_dense.conv.kernel_size,
            stride=self.rbr_dense.conv.stride,
            padding=self.rbr_dense.conv.padding,
            dilation=self.rbr_dense.conv.dilation,
            groups=self.rbr_dense.conv.groups,
            bias=True)
        self.rbr_reparam.weight.data = kernel
        self.rbr_reparam.bias.data = bias
        for para in self.parameters():
            para.detach_()
        self.__delattr__('rbr_dense')
        self.__delattr__('rbr_1x1')
        if hasattr(self, 'rbr_identity'):
            self.__delattr__('rbr_identity')
        if hasattr(self, 'id_tensor'):
            self.__delattr__('id_tensor')
        self.deploy = True


class SimpleEncoder(nn.Module):
    def __init__(self, pretrained=False, rep=False, deploy=False):
        super(SimpleEncoder, self).__init__()
        conv = RepConvBNReLU if rep else BaseConvBNReLU
        self.in_channels = [32, 64, 128, 256]
        self.layer0 = conv(3, self.in_channels[0], 3, 2, 1)
        self.layer1 = nn.Sequential(
            conv(self.in_channels[0], self.in_channels[0], 3, 1, 1, deploy=deploy),
            conv(self.in_channels[0], self.in_channels[0], 3, 1, 1, deploy=deploy),
        )
        self.layer2 = nn.Sequential(
            conv(self.in_channels[0], self.in_channels[1], 3, 2, 1, deploy=deploy),
            conv(self.in_channels[1], self.in_channels[1], 3, 1, 1, deploy=deploy),
        )
        self.layer3 = nn.Sequential(
            conv(self.in_channels[1], self.in_channels[2], 3, 2, 1, deploy=deploy),
            conv(self.in_channels[2], self.in_channels[2], 3, 1, 1, deploy=deploy),
        )
        self.layer4 = nn.Sequential(
            conv(self.in_channels[2], self.in_channels[3], 3, 2, 1, deploy=deploy),
            conv(self.in_channels[3], self.in_channels[3], 3, 1, 1, deploy=deploy),
        )

    def forward(self, x):
        x = self.layer0(x)
        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        x3 = self.layer3(x2)
        x4 = self.layer4(x3)
        return [x1, x2, x3, x4]


class GraduationASPP(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GraduationASPP, self).__init__()

        rate1, rate2, rate3 = (6, 12, 18)

        self.conv_list = nn.ModuleList([
            BaseConvBNReLU(in_channels, out_channels, 1),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate1, rate1),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate2, rate2),
            BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, rate3, rate3),
            nn.Sequential(
                BaseConvBNReLU(in_channels, out_channels // 2, 3, 1, 1),
                BaseConvBNReLU(out_channels // 2, out_channels // 2, 5, 1, 2)
            ),
        ])
        self.pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            BaseConvBNReLU(in_channels, out_channels * 3, 1),
        )

        self.project = BaseConvBNReLU(3 * out_channels, out_channels, 1)

    def forward(self, x):
        size = x.size()[2:]
        res = []
        for conv in self.conv_list:
            res.append(F.interpolate(conv(x), size=size, mode='bilinear', align_corners=False))
        res = torch.cat(res, dim=1)
        pool = base_upsample(self.pool(x), size)
        return self.project(res + pool)


class SpatialAttention(nn.Module):
    def __init__(self):
        super(SpatialAttention, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, 9, padding=4, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.gamma = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avg_out, max_out], dim=1)
        out = self.conv1(out)
        out = self.gamma * out
        return out


class GraduationDetailExtraction(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GraduationDetailExtraction, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(out_channels)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
            nn.BatchNorm2d(out_channels)
        )
        self.relu = nn.ReLU(True)
        self.conv3 = BaseConvBNReLU(in_channels, out_channels, 3, 1, 1)

    def forward(self, x):
        f1 = self.conv1(x)
        f2 = self.conv2(x)
        out = self.relu(f1 + f2 + x)
        out = self.conv3(out)
        return out


class GraduationMiddle(nn.Module):
    def __init__(self, x1_channels, x4_channels, fine_channels):
        super(GraduationMiddle, self).__init__()
        self.aspp = GraduationASPP(x4_channels, x4_channels)
        self.detail_extraction = GraduationDetailExtraction(x1_channels, fine_channels)
        self.down = nn.Sequential(
            BaseConvBNReLU(x4_channels, fine_channels, 3, 1, 1),
            nn.UpsamplingBilinear2d(scale_factor=2),
            BaseConvBNReLU(fine_channels, fine_channels, 3, 1, 1)
        )
        self.edge = BaseConvBNReLU(fine_channels * 2, 1, 1, 1, 0)

    def forward(self, x1, x4):
        f1 = self.detail_extraction(x1)
        f4 = self.aspp(x4)
        down = self.down(f4)
        edge = self.edge(torch.cat([base_upsample(f1, down.size()[2:]), down], dim=1))
        return [f1, f4, edge]


class GraduationOCR(nn.Module):
    def __init__(self, in_channels, out_channels, rep, deploy):
        super(GraduationOCR, self).__init__()
        mid_channels = out_channels // 2

        conv1x1 = RepConv1x1 if rep else BaseConvBNReLU
        conv3x3 = RepConvBNReLU if rep else BaseConvBNReLU

        self.conv = conv3x3(in_channels, mid_channels, 3, 1, 1, deploy=deploy)

        self.f_pixel = conv1x1(mid_channels, mid_channels, 1, deploy=deploy)
        self.f_object = conv1x1(mid_channels, mid_channels, 1, deploy=deploy)
        self.f_down = conv1x1(mid_channels, mid_channels, 1, deploy=deploy)
        self.f_up = conv1x1(mid_channels, in_channels, 1, deploy=deploy)

        self.recover = conv3x3(in_channels * 2, out_channels, 3, 1, 1, deploy=deploy)

    def ocr_gather(self, feat, prob):
        b, c, h, w = feat.size()
        prob = torch.sigmoid(prob).view(b, prob.size(1), -1)  # [b, 1, h*w]
        feat = feat.view(b, c, -1).permute(0, 2, 1)  # [b, h*w, c]
        context = torch.matmul(prob, feat).permute(0, 2, 1).unsqueeze(3)  # [b, c, 1]
        return context

    def ocr_distribute(self, feat, context):
        b, c, h, w = feat.size()
        query = self.f_pixel(feat).view(b, c, -1).permute(0, 2, 1)  # [b, h*w, c]
        key = self.f_object(context).view(b, c, -1)  # [b, c, num_classes]
        value = self.f_down(context).view(b, c, -1).permute(0, 2, 1)  # [b, num_classes, c]
        sim_map = torch.sigmoid(torch.matmul(query, key))  # [b, h*w, num_classes]
        proxy = torch.matmul(sim_map, value)  # [b, h*w, c]
        proxy = proxy.permute(0, 2, 1).contiguous().view(b, c, h, w)
        proxy = self.f_up(proxy)
        return proxy

    def forward(self, x, prob):
        feat = self.conv(x)
        context = self.ocr_gather(feat, prob)
        proxy = self.ocr_distribute(feat, context)
        out = self.recover(torch.cat([proxy, x], dim=1))
        return out


class GraduationDecoder(nn.Module):
    def __init__(self, in_channels, fine_channels, num_classes, rep, deploy):
        super(GraduationDecoder, self).__init__()

        conv = RepConvBNReLU if rep else BaseConvBNReLU

        self.de_en_conv = conv(in_channels * 3, in_channels, 3, 1, 1, deploy=deploy)
        self.sup = nn.Conv2d(in_channels, num_classes, 3, 1, 1)

        self.gate1 = GatedConvBNReLU(in_channels, in_channels, 3, 1, 1, deploy=deploy)
        self.edge = nn.Sequential(
            nn.UpsamplingBilinear2d(scale_factor=2),
            nn.Conv2d(in_channels, 1, 1, 1, 0)
        )

        self.dot = conv(in_channels, fine_channels, 3, 1, 1, deploy=deploy)
        self.dot2 = conv(fine_channels, fine_channels, 3, 1, 1, deploy=deploy)
        self.gate2 = GatedConvBNReLU(fine_channels, fine_channels, 3, 1, 1, deploy=deploy)

        self.conv2 = conv(in_channels + fine_channels, in_channels, 3, 1, 1, deploy=deploy)
        # self.ocr = GraduationOCR(in_channels, in_channels)
        self.conv3 = conv(in_channels, in_channels, 3, 1, 1, deploy=deploy)

    def forward(self, en, de, fine, edge):
        de = base_upsample(de, en.size()[2:])
        de_en = self.de_en_conv(torch.cat([de, en], dim=1))
        sup = self.sup(de_en)

        g1 = self.gate1(en, edge)
        new_edge = self.edge(g1)

        dot = fine * base_upsample(self.dot(de_en), fine.size()[2:])
        dot = self.dot2(dot)
        new_fine = self.gate2(dot, base_upsample(new_edge, dot.size()[2:]))

        pre = self.conv2(torch.cat([de_en, base_upsample(dot, de_en.size()[2:])], dim=1))
        new_de = self.conv3(pre)

        return [new_fine, new_de, new_edge, sup]


class GraduationMiddle2(nn.Module):
    def __init__(self, x1_channels, x4_channels, fine_channels):
        super(GraduationMiddle2, self).__init__()
        self.aspp = GraduationASPP(x4_channels, x4_channels)
        self.detail_extraction = GraduationDetailExtraction(x1_channels, fine_channels)
        self.down = nn.Sequential(
            BaseConvBNReLU(x4_channels, fine_channels, 3, 1, 1),
            nn.UpsamplingBilinear2d(scale_factor=2),
            BaseConvBNReLU(fine_channels, fine_channels, 3, 1, 1)
        )
        self.edge = nn.Conv2d(fine_channels * 2, 1, 1, 1, 0)

    def forward(self, x1, x4):
        f1 = self.detail_extraction(x1)
        f4 = self.aspp(x4)
        down = self.down(f4)
        edge = self.edge(torch.cat([base_upsample(f1, down.size()[2:]), down], dim=1))
        return [f1, f4, edge]


class Module3(nn.Module):
    def __init__(self, in_channels, fine_channels, num_classes, r3, deploy):
        super(Module3, self).__init__()
        conv = RepConvBNReLU if r3 else BaseConvBNReLU
        self.conv1 = conv(fine_channels, fine_channels, 3, 1, 1, deploy=deploy)
        self.gate = GatedConvBNReLU(fine_channels, fine_channels, 3, 1, 1, rep=r3, deploy=deploy)
        self.last_edge = nn.Conv2d(fine_channels, 1, 1, 1, 0)
        self.conv2 = nn.Conv2d(fine_channels + in_channels, num_classes, 3, 1, 1)

    def forward(self, seg, fine, edge):
        fine = self.conv1(fine)
        gate = self.gate(base_upsample(fine, edge.size()[2:]), edge)
        edge = self.last_edge(gate)
        out = self.conv2(torch.cat([base_upsample(seg, edge.size()[2:]), gate], dim=1))
        return out, edge


class Graduation(nn.Module):
    def __init__(self, num_classes, backbone, pretrained=None, rep=False, deploy=False):
        super(Graduation, self).__init__()
        self.backbone = backbone(pretrained=pretrained)
        in_channels = self.backbone.in_channels
        fine_channels = in_channels[0]
        self.middle = GraduationMiddle2(in_channels[0], in_channels[3], fine_channels)
        self.decoder3 = GraduationDecoder(in_channels[2], fine_channels, num_classes, rep, deploy)
        self.decoder2 = GraduationDecoder(in_channels[1], fine_channels, num_classes, rep, deploy)
        self.decoder1 = GraduationDecoder(in_channels[0], fine_channels, num_classes, rep, deploy)
        self.out = Module3(in_channels[0], fine_channels, num_classes, r3=False, deploy=False)

    def forward(self, x):
        x1, x2, x3, x4 = self.backbone(x)
        fine, up, edge1 = self.middle(x1, x4)
        feat1 = fine
        feat2 = up
        fine, up, edge2, sup3 = self.decoder3(x3, up, fine, edge1)
        fine, up, edge3, sup2 = self.decoder2(x2, up, fine, edge2)
        fine, up, edge4, sup1 = self.decoder1(x1, up, fine, edge3)
        out, edge5 = self.out(up, fine, edge4)
        return [out, feat1, feat2]


def graduation_model(num_classes=1, backbone=34, pretrained=None, rep=False, deploy=False):
    if backbone == 18:
        return Graduation(num_classes, ResNet18WithoutMaxPool, pretrained=pretrained, rep=rep, deploy=deploy)
    elif backbone == 34:
        return Graduation(num_classes, ResNet34WithoutMaxPool, pretrained=pretrained, rep=rep, deploy=deploy)
    elif backbone == 10:
        return Graduation(num_classes, ResNet10WithoutMaxPool, pretrained=pretrained, rep=rep, deploy=deploy)
    elif backbone == 0:
        return Graduation(num_classes, SimpleEncoder, pretrained=pretrained, rep=rep, deploy=deploy)
